import pandas as pd
import io
from contextlib import redirect_stdout, redirect_stderr
import traceback
import ast
import streamlit as st
import numpy as np
import math
import random
import json
import re
import time
import copy


class CodeInterpreter:
    def __init__(self):
        self.allowed_modules = {
            'pandas': pd,
            'numpy': np,
            'math': math,
            'random': random,
            'json': json,
            're': re,
            'time': time,
            'copy': copy
        }

    def is_safe_code(self, code: str) -> bool:
        """检查代码是否安全可执行"""
        try:
            tree = ast.parse(code)
            for node in ast.walk(tree):
                # 检查是否有危险的内置函数调用
                if isinstance(node, ast.Call):
                    if isinstance(node.func, ast.Name):
                        if node.func.id in ['eval', 'exec', 'compile', '__import__']:
                            return False
                # 检查导入语句
                elif isinstance(node, (ast.Import, ast.ImportFrom)):
                    for name in node.names:
                        module = name.name.split('.')[0]
                        if module not in self.allowed_modules:
                            return False
            return True
        except SyntaxError:
            return False

    def prepare_environment(self, df=None):
        """准备代码执行环境"""
        env = {
            'pd': pd,
            'df': df,
            'st': st,
            'np': np,
            'math': math,
            'random': random,
            'json': json,
            're': re,
            'time': time,
            'copy': copy
        }
        return env

    def execute_code(self, code: str, df=None):
        """执行代码并返回结果"""
        if not self.is_safe_code(code):
            return {
                'success': False,
                'error': "代码包含不允许的操作或模块",
                'output': None,
                'df': df
            }

        output = io.StringIO()
        local_dict = self.prepare_environment(df)

        try:
            with redirect_stdout(output), redirect_stderr(output):
                exec(code, local_dict)

            return {
                'success': True,
                'error': None,
                'output': output.getvalue(),
                'df': local_dict.get('df', df)
            }
        except Exception as e:
            return {
                'success': False,
                'error': traceback.format_exc(),
                'output': output.getvalue(),
                'df': df
            }

    def format_code(self, code: str) -> str:
        """格式化代码，移除多余的空行和注释"""
        lines = code.split('\n')
        formatted_lines = []
        for line in lines:
            stripped = line.strip()
            if stripped and not stripped.startswith('#'):
                formatted_lines.append(line)
        return '\n'.join(formatted_lines) 